import requests
import time
import argparse

class StreamTest:
    def __init__(self, base_url, token):
        self.base_url = base_url
        self.headers = {
            'Authorization': f'Bearer {token}'
        }

    def upload_stream(self, rtmp_url):
        """
        测试流上传接口
        """
        url = f"{self.base_url}/api/stream/upload"
        data = {
            'stream': rtmp_url,
            'stream_name': f'test_stream_{int(time.time())}'
        }
        
        try:
            response = requests.post(url, params=data, headers=self.headers)
            print(f"Upload Response Status: {response.status_code}")
            print(f"Upload Response Body: {response.json()}")
            return response.json()
        except Exception as e:
            print(f"Upload Error: {str(e)}")
            return None

    def play_stream(self, stream_name):
        """
        测试流播放接口
        """
        url = f"{self.base_url}/api/stream/play"
        params = {
            'stream': stream_name
        }
        
        try:
            response = requests.get(url, params=params, headers=self.headers, stream=True)
            print(f"Play Response Status: {response.status_code}")
            
            # 如果是流媒体响应，打印前几个字节的内容
            if response.status_code == 200:
                print("Receiving stream data...")
                chunk = next(response.iter_content(chunk_size=1024))
                print(f"First chunk size: {len(chunk)} bytes")
            else:
                print(f"Play Response Body: {response.json()}")
                
        except Exception as e:
            print(f"Play Error: {str(e)}")

def main():
    parser = argparse.ArgumentParser(description='Test Stream API')
    parser.add_argument('--url', default='http://10.8.49.64:9091', help='Base URL of the API')
    parser.add_argument('--token', required=True, help='a913519db8588229fa7cca43037089556771f294')
    parser.add_argument('--stream', required=True, help='rtmp://172.31.100.201/meeting/922b0edb8a944264befbca5d9e10b3f2?streamToken=8678396fcc2c417cbc687eb2f85a044e')
    
    args = parser.parse_args()
    
    tester = StreamTest(args.url, args.token)
    
    # 测试上传
    print("Testing stream upload...")
    result = tester.upload_stream(args.stream)
    
    if result and result.get('code') == 200:
        # 如果上传成功，测试播放
        print("\nTesting stream play...")
        stream_name = result.get('data', {}).get('stream_name')
        if stream_name:
            tester.play_stream(stream_name)

if __name__ == '__main__':
    main()